# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

# -*- coding: utf-8 -*-
"""
测试字典方法的代码生成

验证 dict 类型的方法能够正确生成 C++ 代码，
特别关注视图对象、临时对象生命周期、弱类型处理等关键机制。
"""

from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains
from typing import Any


def test_dict_get_method():
    """测试 dict.get() 方法的代码生成"""

    def test_get(d: dict[str, int]) -> int:
        """测试 get 方法"""
        return d.get("key", 0)

    cpp = compile_chunk(test_get)
    expected = """
    int test_get(const mc::dict& d) {
        return mcpy::dict_get<int>(d, "key", 0);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_keys_view():
    """测试 dict.keys() 视图对象的代码生成"""

    code = """
def test_keys(d: dict[str, int]) -> list[str]:
    '''测试 keys 视图'''
    return list(d.keys())
"""

    cpp = compile_chunk(code)
    expected = """
    mc::array<std::string> test_keys(const mc::dict& d) {
        return mcpy::list_construct<std::string>(mcpy::dict_keys<std::string>(d));
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_values_view():
    """测试 dict.values() 视图对象的代码生成"""

    code = """
def test_values(d: dict[str, int]) -> list[int]:
    '''测试 values 视图'''
    return list(d.values())
"""

    cpp = compile_chunk(code)
    expected = """
    mc::array<int> test_values(const mc::dict& d) {
        return mcpy::list_construct<int>(mcpy::dict_values<int>(d));
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_items_view():
    """测试 dict.items() 视图对象的代码生成"""

    code = """
def test_items(d: dict[str, int]) -> list[tuple[str, int]]:
    '''测试 items 视图'''
    return list(d.items())
"""

    cpp = compile_chunk(code)
    expected = """
    mc::array<std::tuple<std::string, int>> test_items(const mc::dict& d) {
        return mcpy::list_construct<std::tuple<std::string, int>>(mcpy::dict_items<std::string, int>(d));
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_literal_keys_view():
    """测试字典字面量的 keys() 视图（临时对象）"""

    code = """
def test_literal_keys() -> None:
    '''测试字面量字典的 keys 视图 - 验证临时对象生命周期'''
    for key in {"a": 1, "b": 2}.keys():
        pass
"""

    cpp = compile_chunk(code)
    # 关键：应该传递临时对象（rvalue），验证生命周期正确
    # 字面量字典被推导为弱类型 dict，键类型为 std::string（从字面量推导）
    expected = """
    void test_literal_keys() {
        for (const auto& key : mcpy::dict_keys<std::string>(mc::dict{{"a", 1}, {"b", 2}})) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_function_return_views():
    """测试函数返回字典后调用视图方法（最关键的临时对象场景）"""

    code = """
def get_config() -> dict[str, int]:
    '''返回配置字典'''
    return {"host": 127, "port": 8080}

def test_function_return_keys() -> None:
    '''测试函数返回的临时字典调用 keys()'''
    for key in get_config().keys():
        pass
"""

    cpp = compile_chunk(code)
    # 关键：get_config() 返回临时对象，应该正确处理生命周期
    # 返回类型标注为 dict[str, int]，所以 keys() 的类型参数是 std::string
    expected_get_config = """
    mc::dict get_config() {
        return mc::dict{{"host", 127}, {"port", 8080}};
    }
    """
    expected_test = """
    void test_function_return_keys() {
        for (const auto& key : mcpy::dict_keys<std::string>(get_config())) {
        }
    }
    """
    assert_code_contains(cpp, expected_get_config)
    assert_code_contains(cpp, expected_test)


def test_dict_weak_type():
    """测试弱类型 dict 的代码生成"""

    code = """
def test_weak_dict() -> None:
    '''测试弱类型字典'''
    d: dict = {"a": 1, "b": "hello"}
    value = d["a"]
"""

    cpp = compile_chunk(code)
    expected = """
    void test_weak_dict() {
        mc::dict d = mc::dict{{"a", 1}, {"b", "hello"}};
        mc::variant value = d["a"];
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_weak_type_methods():
    """测试弱类型 dict 调用方法"""

    code = """
def get_data() -> dict:
    '''返回弱类型字典'''
    return {"x": 10}

def test_weak_dict_methods() -> None:
    '''测试弱类型字典的方法调用'''
    for key in get_data().keys():
        pass
"""

    cpp = compile_chunk(code)
    # 弱类型 dict 的视图应该使用 mc::variant 作为类型参数
    expected = """
    mc::dict get_data() {
        return mc::dict{{"x", 10}};
    }
    void test_weak_dict_methods() {
        for (const auto& key : mcpy::dict_keys<mc::variant>(get_data())) {
    """
    assert_code_contains(cpp, expected)


def test_dict_pop_method():
    """测试 dict.pop() 方法的代码生成"""

    code = """
def test_pop(d: dict[str, int]) -> int:
    '''测试 pop 方法'''
    return d.pop("key", 999)
"""

    cpp = compile_chunk(code)
    expected = """
    int test_pop(mc::dict& d) {
        return mcpy::dict_pop<int>(d, "key", 999);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_update_method():
    """测试 dict.update() 方法的代码生成"""

    code = """
def test_update(d1: dict[str, int], d2: dict[str, int]) -> None:
    '''测试 update 方法'''
    d1.update(d2)
"""

    cpp = compile_chunk(code)
    expected = """
    void test_update(mc::dict& d1, const mc::dict& d2) {
        mcpy::dict_update(d1, d2);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_setdefault_method():
    """测试 dict.setdefault() 方法的代码生成"""

    code = """
def test_setdefault(d: dict[str, int]) -> None:
    '''测试 setdefault 方法'''
    d.setdefault("key", 100)
"""

    cpp = compile_chunk(code)
    expected = """
    void test_setdefault(mc::dict& d) {
        mcpy::dict_setdefault(d, "key", 100);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_merge_operator():
    """测试 dict | dict 合并操作符的代码生成"""

    code = """
def test_merge(d1: dict[str, int], d2: dict[str, int]) -> dict[str, int]:
    '''测试 | 合并操作符'''
    return d1 | d2
"""

    cpp = compile_chunk(code)
    expected = """
    mc::dict test_merge(const mc::dict& d1, const mc::dict& d2) {
        return mcpy::dict_merge(d1, d2);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_inplace_merge_operator():
    """测试 dict |= dict 就地合并操作符的代码生成"""

    code = """
def test_inplace_merge(d1: dict[str, int], d2: dict[str, int]) -> None:
    '''测试 |= 就地合并操作符'''
    d1 |= d2
"""

    cpp = compile_chunk(code)
    expected = """
    void test_inplace_merge(mc::dict& d1, const mc::dict& d2) {
        mcpy::dict_ior(d1, d2);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_fromkeys_classmethod():
    """测试 dict.fromkeys() 类方法的代码生成"""

    code = """
def test_fromkeys() -> dict[str, int]:
    '''测试 fromkeys 类方法'''
    keys = ["a", "b", "c"]
    return dict.fromkeys(keys, 0)
"""

    cpp = compile_chunk(code)
    expected = """
    mc::dict test_fromkeys() {
        mc::array<std::string> keys = {"a", "b", "c"};
        return mcpy::dict_fromkeys(keys, 0);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_reversed_function():
    """测试 reversed(dict) 全局函数的代码生成"""

    code = """
def test_reversed_dict(d: dict[str, int]) -> None:
    '''测试反向迭代字典'''
    for key in reversed(d):
        pass
"""

    cpp = compile_chunk(code)
    expected = """
    void test_reversed_dict(const mc::dict& d) {
        for (const auto& key : mcpy::dict_reversed(d)) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_reversed_literal():
    """测试字面量字典的反向迭代（临时对象）"""

    code = """
def test_reversed_literal() -> None:
    '''测试字面量字典的反向迭代'''
    for key in reversed({"a": 1, "b": 2, "c": 3}):
        pass
"""

    cpp = compile_chunk(code)
    # 应该使用 rvalue 重载，传递临时对象
    expected = """
    void test_reversed_literal() {
        for (const auto& key : mcpy::dict_reversed(mc::dict{{"a", 1}, {"b", 2}, {"c", 3}})) {
        }
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_nested_access():
    """测试嵌套字典访问的代码生成"""

    code = """
def test_nested(d: dict) -> int:
    '''测试嵌套字典访问'''
    return d["a"]["b"]["c"]
"""

    cpp = compile_chunk(code)
    # 链式访问应该生成正确的代码
    assert_code_contains(cpp, 'return mcpy::cast<int>(d["a"]["b"]["c"]);')


def test_dict_chain_assignment():
    """测试字典链式赋值的代码生成"""

    code = """
def test_chain_assign(d: dict) -> None:
    '''测试链式赋值'''
    d["user"]["name"] = "Alice"
"""

    cpp = compile_chunk(code)
    # 链式赋值应该生成 mutable 引用
    assert_code_contains(cpp, 'd["user"]["name"] = "Alice";')


# ============================================================================
# dict 构造函数测试
# ============================================================================


def test_dict_empty_constructor():
    """测试空 dict 构造函数"""

    def test_empty() -> None:
        """测试空构造"""
        d1: dict[str, int] = dict()
        d2: dict[int, str] = dict()

    cpp = compile_chunk(test_empty)
    expected = """
    void test_empty() {
        mc::dict d1 = mc::dict();
        mc::dict d2 = mc::dict();
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_from_list_constructor():
    """测试从键值对列表构造 dict"""

    def test_from_list() -> None:
        """从列表构造 dict"""
        items: list[tuple[str, int]] = [("a", 1), ("b", 2), ("c", 3)]
        d: dict[str, int] = dict(items)

    cpp = compile_chunk(test_from_list)
    expected = """
    void test_from_list() {
        mc::array<std::tuple<std::string, int>> items = {std::make_tuple("a", 1), std::make_tuple("b", 2), std::make_tuple("c", 3)};
        mc::dict                                  d     = mcpy::dict_construct(items);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_from_literal_list_constructor():
    """测试从列表字面量构造 dict"""

    def test_from_literal() -> None:
        """从列表字面量构造 dict"""
        d: dict[str, int] = dict([("a", 1), ("b", 2)])

    cpp = compile_chunk(test_from_literal)
    expected = """
    void test_from_literal() {
        mc::dict d = mcpy::dict_construct(mc::array<std::tuple<std::string, int>>{std::make_tuple("a", 1), std::make_tuple("b", 2)});
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_from_dict_constructor():
    """测试从另一个 dict 拷贝构造"""

    def test_copy_constructor(d1: dict[str, int]) -> dict[str, int]:
        """拷贝构造 dict"""
        d2: dict[str, int] = dict(d1)
        return d2

    cpp = compile_chunk(test_copy_constructor)
    expected = """
    mc::dict test_copy_constructor(const mc::dict& d1) {
        mc::dict d2 = mc::dict(d1);
        return d2;
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_constructor_type_inference():
    """测试 dict 构造函数的类型推导"""

    def test_inference() -> None:
        """测试类型推导"""
        # 带类型注解
        d1: dict[str, int] = dict([("a", 1), ("b", 2)])
        # 从已知类型列表推导
        items: list[tuple[str, int]] = [("x", 10), ("y", 20)]
        d2: dict[str, int] = dict(items)

    cpp = compile_chunk(test_inference)
    expected = """
    void test_inference() {
        mc::dict                                  d1    = mcpy::dict_construct(mc::array<std::tuple<std::string, int>>{std::make_tuple("a", 1), std::make_tuple("b", 2)});
        mc::array<std::tuple<std::string, int>> items = {std::make_tuple("x", 10), std::make_tuple("y", 20)};
        mc::dict                                  d2    = mcpy::dict_construct(items);
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_constructor_in_expression():
    """测试 dict 构造函数在表达式中使用"""

    def test_in_expr() -> int:
        """在表达式中使用 dict 构造"""
        return len(dict([("a", 1), ("b", 2), ("c", 3)]))

    cpp = compile_chunk(test_in_expr)
    expected = """
    int test_in_expr() {
        return mcpy::dict_construct(mc::array<std::tuple<std::string, int>>{std::make_tuple("a", 1), std::make_tuple("b", 2), std::make_tuple("c", 3)}).size();
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_empty_weak_type_constructor():
    """测试弱类型空 dict 构造"""

    def test_weak_empty() -> None:
        """弱类型空 dict"""
        d: dict = dict()

    cpp = compile_chunk(test_weak_empty)
    expected = """
    void test_weak_empty() {
        mc::dict d = mc::dict();
    }
    """
    assert_code_contains(cpp, expected)


def test_dict_for_in_loop():
    """验证 python 字典迭代的是 keys() 视图"""

    def for_in_loop() -> None:
        d = {"a": 1, "b": 2, "c": 3, 2: 1}
        count = 0

        for key in d:
            count += d[key]

        count = 0
        for key in d.keys():
            count += d[key]

    cpp = compile_chunk(for_in_loop)
    expected = """
    void for_in_loop() {
        mc::dict d = mc::dict{{"a", 1}, {"b", 2}, {"c", 3}, {2, 1}};
        int count = 0;
        for (const auto& key : mcpy::dict_keys<mc::variant>(d)) {
            count += mcpy::cast<int>(d[key]);
        }

        count = 0;
        for (const auto& key : mcpy::dict_keys<mc::variant>(d)) {
            count += mcpy::cast<int>(d[key]);
        }
    }
    """
    assert_code_contains(cpp, expected)
